package hn.cch.spring_boot_security_rest.wrapper;

import hn.cch.spring_boot_security_rest.constant.SystemInfo;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

public class ApiResponseWrapper extends HttpServletResponseWrapper {



    private ByteArrayOutputStream baos = null;

    private ServletOutputStream sos = null;

    private PrintWriter pw = null;


    public ApiResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        baos = new ByteArrayOutputStream();
        sos = new ApiServletOutputStream(baos);
        pw = new PrintWriter(new OutputStreamWriter(baos, SystemInfo.CHARSET));
    }


    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return sos;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return pw;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (sos != null) {
            sos.flush();
        }
        if (pw != null) {
            pw.flush();
        }
    }

    @Override
    public void reset() {
        baos.reset();
    }



    public byte[] getResponseBody() throws IOException {
        //将sos、pw中的数据强制输出到baos里面，否则取不到数据
        flushBuffer();
        return baos.toByteArray();
    }


}
